import torch
import copy
import torch.nn as nn
from torch.nn import functional as F
from torch.nn import Module
from torch.nn import MultiheadAttention
from torch.nn import ModuleList
from torch.nn.init import xavier_uniform_
from torch.nn import Dropout
from torch.nn import Linear
from torch.nn import LayerNorm


class Encoder(Module):
    r"""Transformer编码器栈：由N个编码器层堆叠组成
    用于提取序列数据的特征，捕捉序列内部的依赖关系（如时序、空间关系）

    Args:
        encoder_layer: 单个编码器层实例（EncoderLayer类对象，必需）
        num_layers: 编码器层的堆叠数量（必需）
        norm: 输出的层归一化组件（可选，默认None）
    """

    def __init__(self, encoder_layer, num_layers, norm=None):
        super(Encoder, self).__init__()
        # 复制N个编码器层组成层列表（保证各层参数独立）
        self.layers = _get_clones(encoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm  # 可选的输出归一化

    def forward(self, src):
        r"""前向传播：将输入依次通过所有编码器层
        Args:
            src: 输入序列特征 (seq_len, batch_size, d_model)
        Returns:
            经过所有编码器层处理后的特征 (seq_len, batch_size, d_model)
        """
        output = src
        # 逐层传递输入特征
        for i in range(self.num_layers):
            output = self.layers[i](output)
        # 可选的输出归一化
        if self.norm:
            output = self.norm(output)
        return output


class Decoder(Module):
    r"""Transformer解码器栈：由N个解码器层堆叠组成
    用于结合编码器的记忆特征（memory）对目标序列（tgt）进行解码，实现跨序列依赖捕捉

    Args:
        decoder_layer: 单个解码器层实例（DecoderLayer类对象，必需）
        num_layers: 解码器层的堆叠数量（必需）
        norm: 输出的层归一化组件（可选，默认None）
    """

    def __init__(self, decoder_layer, num_layers, norm=None):
        super(Decoder, self).__init__()
        # 复制N个解码器层组成层列表
        self.layers = _get_clones(decoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm  # 可选的输出归一化

    def forward(self, tgt, memory):
        r"""前向传播：将目标序列和记忆特征依次通过所有解码器层
        Args:
            tgt: 目标序列特征 (tgt_seq_len, batch_size, d_model)
            memory: 编码器输出的记忆特征 (src_seq_len, batch_size, d_model)
        Returns:
            经过所有解码器层处理后的输出特征 (tgt_seq_len, batch_size, d_model)
        """
        output = tgt
        # 逐层传递目标序列和记忆特征
        for i in range(self.num_layers):
            output = self.layers[i](output, memory)
        # 可选的输出归一化
        if self.norm:
            output = self.norm(output)
        return output


class EncoderLayer(Module):
    r"""单个Transformer编码器层（借鉴自CMRAN模型）
    核心结构：自注意力机制 + 前馈神经网络，配合残差连接和层归一化

    Args:
        d_model: 输入特征的维度（必需）
        nhead: 多头注意力机制的头数（必需）
        dim_feedforward: 前馈神经网络的中间层维度（默认1024）
        dropout: Dropout概率（默认0.1）
        activation: 前馈网络中间层的激活函数（默认relu，可选gelu）
    """

    def __init__(self, d_model, nhead, dim_feedforward=1024, dropout=0.1, activation="relu"):
        super(EncoderLayer, self).__init__()
        # 自注意力层：捕捉序列内部的依赖关系
        self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
        # 前馈神经网络：由两个线性层和Dropout组成
        self.linear1 = Linear(d_model, dim_feedforward)  # 升维
        self.dropout = Dropout(dropout)
        self.linear2 = Linear(dim_feedforward, d_model)  # 降维

        # 层归一化层：稳定训练过程
        self.norm1 = LayerNorm(d_model)
        self.norm2 = LayerNorm(d_model)
        # Dropout层：防止过拟合
        self.dropout1 = Dropout(dropout)
        self.dropout2 = Dropout(dropout)

        # 激活函数：根据参数选择relu或gelu
        self.activation = _get_activation_fn(activation)

    def forward(self, src):
        r"""前向传播：执行编码器层的核心计算
        Args:
            src: 输入序列特征 (seq_len, batch_size, d_model)
        Returns:
            经过编码器层处理后的特征 (seq_len, batch_size, d_model)
        """
        # 1. 自注意力计算 + 残差连接 + Dropout + 层归一化
        src2 = self.self_attn(src, src, src)[0]  # 自注意力输出（q=k=v=src）
        src = src + self.dropout1(src2)  # 残差连接 + Dropout
        src = self.norm1(src)  # 层归一化

        # 2. 前馈网络计算 + 残差连接 + Dropout + 层归一化
        if hasattr(self, "activation"):
            src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
        else:  # 向后兼容：默认使用relu
            src2 = self.linear2(self.dropout(F.relu(self.linear1(src))))
        src = src + self.dropout2(src2)  # 残差连接 + Dropout
        src = self.norm2(src)  # 层归一化

        return src


class DecoderLayer(Module):
    r"""单个Transformer解码器层（借鉴自CMRAN模型）
    核心结构：多头交叉注意力机制 + 前馈神经网络，配合残差连接和层归一化
    注：与标准Transformer解码器不同，该层未包含自注意力层，专注于跨模态交互

    Args:
        d_model: 输入特征的维度（必需）
        nhead: 多头注意力机制的头数（必需）
        dim_feedforward: 前馈神经网络的中间层维度（默认1024）
        dropout: Dropout概率（默认0.1）
        activation: 前馈网络中间层的激活函数（默认relu，可选gelu）
    """

    def __init__(self, d_model, nhead, dim_feedforward=1024, dropout=0.1, activation="relu"):
        super(DecoderLayer, self).__init__()
        # 自注意力层（当前实现中未使用，保留结构）
        self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
        # 多头交叉注意力层：捕捉目标序列与记忆特征的依赖关系
        self.multihead_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
        # 前馈神经网络：与编码器层结构一致
        self.linear1 = Linear(d_model, dim_feedforward)
        self.dropout = Dropout(dropout)
        self.linear2 = Linear(dim_feedforward, d_model)

        # 层归一化层
        self.norm1 = LayerNorm(d_model)
        self.norm2 = LayerNorm(d_model)
        # Dropout层
        self.dropout1 = Dropout(dropout)
        self.dropout2 = Dropout(dropout)

        # 激活函数
        self.activation = _get_activation_fn(activation)

    def forward(self, tgt, memory):
        r"""前向传播：执行解码器层的核心计算
        Args:
            tgt: 目标序列特征 (tgt_seq_len, batch_size, d_model)
            memory: 编码器输出的记忆特征 (src_seq_len, batch_size, d_model)
        Returns:
            经过解码器层处理后的输出特征 (tgt_seq_len, batch_size, d_model)
        """
        # 拼接记忆特征和目标序列，扩展注意力的上下文信息
        memory = torch.cat([memory, tgt], dim=0)
        # 1. 交叉注意力计算 + 残差连接 + Dropout + 层归一化
        tgt2 = self.multihead_attn(tgt, memory, memory)[0]  # 交叉注意力输出（q=tgt, k=v=memory）
        tgt = tgt + self.dropout1(tgt2)  # 残差连接 + Dropout
        tgt = self.norm1(tgt)  # 层归一化

        # 2. 前馈网络计算 + 残差连接 + Dropout + 层归一化
        if hasattr(self, "activation"):
            tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
        else:  # 向后兼容
            tgt2 = self.linear2(self.dropout(F.relu(self.linear1(tgt))))
        tgt = tgt + self.dropout2(tgt2)  # 残差连接 + Dropout
        tgt = self.norm2(tgt)  # 层归一化

        return tgt


def _get_clones(module, N):
    r"""辅助函数：复制N个相同的模块，组成ModuleList（保证各模块参数独立）
    Args:
        module: 待复制的模块
        N: 复制次数
    Returns:
        ModuleList: 包含N个复制模块的列表
    """
    return ModuleList([copy.deepcopy(module) for i in range(N)])


def _get_activation_fn(activation):
    r"""辅助函数：根据字符串参数返回对应的激活函数
    Args:
        activation: 激活函数名称（"relu"或"gelu"）
    Returns:
        对应的激活函数
    Raises:
        RuntimeError: 激活函数名称不合法时抛出异常
    """
    if activation == "relu":
        return F.relu
    elif activation == "gelu":
        return F.gelu
    else:
        raise RuntimeError("activation should be relu/gelu, not %s." % activation)


class New_Audio_Guided_Attention(nn.Module):
    r"""音频引导的注意力模块：融合音频特征对视频特征进行空间和通道维度的增强
    注：实际实现与论文描述略有差异，经优化后效率更高
    核心功能：通过音频特征引导，生成视频的通道注意力权重和空间注意力权重，同时结合自注意力进一步优化视频特征
    """

    def __init__(self, beta):
        super(New_Audio_Guided_Attention, self).__init__()

        self.beta = beta  # 自注意力特征的融合权重
        self.relu = nn.ReLU()  # ReLU激活函数
        self.video_input_dim = 512  # 视频输入特征维度
        self.audio_input_dim = 128  # 音频输入特征维度
        self.hidden_dim = 256  # 中间隐藏层维度

        # 1. 通道注意力（Channel Attention）相关层：调节视频特征的通道重要性
        self.affine_video_1 = nn.Linear(self.video_input_dim, self.video_input_dim)  # 视频特征线性映射
        self.affine_audio_1 = nn.Linear(self.audio_input_dim, self.video_input_dim)  # 音频特征线性映射（适配视频维度）
        self.affine_bottleneck = nn.Linear(self.video_input_dim, self.hidden_dim)  # 瓶颈层降维
        self.affine_v_c_att = nn.Linear(self.hidden_dim, self.video_input_dim)  # 生成通道注意力权重

        # 2. 音频引导的空间注意力（Spatial Attention）相关层：调节视频特征的空间重要性
        self.affine_video_2 = nn.Linear(self.video_input_dim, self.hidden_dim)  # 视频特征线性映射
        self.affine_audio_2 = nn.Linear(self.audio_input_dim, self.hidden_dim)  # 音频特征线性映射
        self.affine_v_s_att = nn.Linear(self.hidden_dim, 1)  # 生成空间注意力权重

        # 3. 视频自注意力（Self Attention）相关层：捕捉视频内部的空间依赖
        self.latent_dim = 4  # latent维度（用于降维计算自注意力）
        self.video_query = nn.Linear(self.video_input_dim, self.video_input_dim // self.latent_dim)  # 查询向量生成
        self.video_key = nn.Linear(self.video_input_dim, self.video_input_dim // self.latent_dim)  # 键向量生成
        self.video_value = nn.Linear(self.video_input_dim, self.video_input_dim)  # 值向量生成

        # 4. 视频自空间注意力相关层：基于视频全局平均特征的空间注意力
        self.affine_video_ave = nn.Linear(self.video_input_dim, self.hidden_dim)  # 全局平均特征映射
        self.affine_video_3 = nn.Linear(self.video_input_dim, self.hidden_dim)  # 视频特征映射
        self.ave_bottleneck = nn.Linear(512, 256)  # 瓶颈层（当前未使用）
        self.ave_v_att = nn.Linear(self.hidden_dim, 1)  # 生成自空间注意力权重

        self.tanh = nn.Tanh()  # Tanh激活函数
        self.softmax = nn.Softmax(dim=-1)  # Softmax归一化（用于注意力权重）
        self.dropout = nn.Dropout(0.2)  # Dropout层防止过拟合
        self.norm = nn.LayerNorm(self.video_input_dim)  # 层归一化（稳定自注意力输出）

    def forward(self, video, audio):
        '''
        前向传播：融合音频特征生成增强的视频特征
        Args:
            video: 视频特征 (batch_size, seq_len, h, w, video_input_dim) → 这里h=7, w=7（空间维度）
            audio: 音频特征 (seq_len, batch_size, audio_input_dim)
        Returns:
            c_s_att_visual_feat: 增强后的视频特征 (batch_size, seq_len, video_input_dim)
        '''
        # 调整音频特征维度：(seq_len, batch_size, a_dim) → (batch_size, seq_len, a_dim)
        audio = audio.transpose(1, 0)
        # 获取输入特征的维度信息
        batch, t_size, h, w, v_dim = video.size()  # batch：批次大小，t_size：时间步长，h/w：空间尺寸，v_dim：视频维度
        a_dim = audio.size(-1)  # 音频维度
        # 音频特征reshape：(batch, t_size, a_dim) → (batch*t_size, a_dim)（便于批量计算）
        audio_feature = audio.reshape(batch * t_size, a_dim)
        # 视频特征reshape：(batch, t_size, h, w, v_dim) → (batch, t_size, h*w, v_dim)（展平空间维度）
        visual_feature = video.reshape(batch, t_size, -1, v_dim)
        raw_visual_feature = visual_feature  # 保存原始视频特征（用于后续残差连接）

        # ============================== 视频自注意力（Video Self Attention）=======================================
        # 生成查询、键、值向量并调整维度
        video_query_feature = self.video_query(visual_feature).reshape(batch * t_size, h * w,
                                                                       -1)  # (batch*t_size, h*w, v_dim/latent_dim)
        video_key_feature = self.video_key(visual_feature).reshape(batch * t_size, h * w, -1).permute(0, 2,
                                                                                                      1)  # (batch*t_size, v_dim/latent_dim, h*w)
        # 计算注意力分数（能量值）并归一化
        energy = torch.bmm(video_query_feature, video_key_feature)  # (batch*t_size, h*w, h*w)
        attention = self.softmax(energy)  # 空间维度归一化
        # 计算自注意力输出并进行残差连接和归一化
        video_value_feature = self.video_value(visual_feature).reshape(batch * t_size, h * w,
                                                                       -1)  # (batch*t_size, h*w, v_dim)
        output = torch.matmul(attention, video_value_feature)  # 自注意力加权输出
        # 残差连接 + Dropout + 层归一化
        output = self.norm(visual_feature.reshape(batch * t_size, h * w, -1) + self.dropout(output))
        visual_feature = output  # 更新视频特征为自注意力增强后的值

        # ============================== 视频自空间注意力（Video Self Spatial Attention）===================================
        # 计算视频特征的全局平均（空间维度）
        video_average = visual_feature.sum(dim=1) / (h * w)  # (batch, t_size, v_dim) → 空间维度平均
        video_average = video_average.reshape(batch * t_size, v_dim)  # (batch*t_size, v_dim)
        video_average = self.relu(self.affine_video_ave(video_average)).unsqueeze(-2)  # (batch*t_size, 1, hidden_dim)
        # 视频特征映射并与全局平均特征交互
        self_video_att_feat = visual_feature.reshape(batch * t_size, -1, v_dim)  # (batch*t_size, h*w, v_dim)
        self_video_att_query = self.relu(self.affine_video_3(self_video_att_feat))  # (batch*t_size, h*w, hidden_dim)
        self_query = self_video_att_query * video_average  # 元素乘交互（融合全局信息）
        # 生成空间注意力权重并加权视频特征
        self_spatial_att_maps = self.softmax(
            self.tanh(self.ave_v_att(self_query)).transpose(2, 1))  # (batch*t_size, 1, h*w)
        self_att_feat = torch.bmm(self_spatial_att_maps, visual_feature).squeeze().reshape(batch, t_size,
                                                                                           v_dim)  # (batch, t_size, v_dim)

        # ============================== 音频引导的通道注意力（Audio-Guided Channel Attention）===================================
        # 音频和视频特征映射并交互
        audio_query_1 = self.relu(self.affine_audio_1(audio_feature)).unsqueeze(-2)  # (batch*t_size, 1, v_dim)
        video_query_1 = self.relu(self.affine_video_1(visual_feature)).reshape(batch * t_size, h * w,
                                                                               -1)  # (batch*t_size, h*w, v_dim)
        audio_video_query_raw = (audio_query_1 * video_query_1).mean(-2)  # 元素乘后空间维度平均（融合音视频信息）
        # 生成通道注意力权重
        audio_video_query = self.relu(self.affine_bottleneck(audio_video_query_raw))  # 瓶颈层降维
        channel_att_maps = self.affine_v_c_att(audio_video_query).sigmoid().reshape(batch, t_size, -1,
                                                                                    v_dim)  # (batch, t_size, h*w, v_dim)
        # 通道注意力加权原始视频特征
        c_att_visual_feat = (raw_visual_feature * (channel_att_maps + 1))  # +1保证权重不低于1，避免特征被过度抑制

        # ============================== 音频引导的空间注意力（Audio-Guided Spatial Attention）====================================
        # 调整通道注意力加权后的视频特征维度
        c_att_visual_feat = c_att_visual_feat.reshape(batch * t_size, -1, v_dim)  # (batch*t_size, h*w, v_dim)
        # 音视频特征映射并交互
        c_att_visual_query = self.relu(self.affine_video_2(c_att_visual_feat))  # (batch*t_size, h*w, hidden_dim)
        audio_query_2 = self.relu(self.affine_audio_2(audio_feature)).unsqueeze(-2)  # (batch*t_size, 1, hidden_dim)
        audio_video_query_2 = c_att_visual_query * audio_query_2  # 元素乘交互
        # 生成空间注意力权重并加权视频特征
        spatial_att_maps = self.softmax(
            self.tanh(self.affine_v_s_att(audio_video_query_2)).transpose(2, 1))  # (batch*t_size, 1, h*w)
        c_s_att_visual_feat = torch.bmm(spatial_att_maps, c_att_visual_feat).squeeze().reshape(batch, t_size,
                                                                                               v_dim)  # (batch, t_size, v_dim)

        # 融合自注意力特征：残差连接（beta控制自注意力的贡献）
        c_s_att_visual_feat = c_s_att_visual_feat + self.beta * self_att_feat.sigmoid() * c_s_att_visual_feat

        return c_s_att_visual_feat